import json
import logging
import os
import time
import pandas as pd
import baostock as bs
import mplfinance as mpf
import adata

from Indicators import Boll
from Indicators.macd import *
from main import get_current_history_candles
from tools import picture_upload
from tools.dingdingrobot import call_dingtalk
from tools.send_mail import send_email


def send_mail_and_message(df, industry, code,name = None):
    # 获取布林带算过的history_candles
    bolling_candles = Boll.Boll(df, 20, 2)
    if  "date" in bolling_candles.columns:
        bolling_candles.set_index(pd.to_datetime(bolling_candles['date']),inplace=True)
    elif "trade_date" in bolling_candles.columns:
        bolling_candles.set_index(pd.to_datetime(bolling_candles['trade_date']),inplace=True)

    # 创建Bollinger Bands的添加线
    apds = [
        mpf.make_addplot(bolling_candles['MiddleBand'], color='blue'),
        mpf.make_addplot(bolling_candles['UpperBand'], color='red'),
        mpf.make_addplot(bolling_candles['LowerBand'], color='green'),
        mpf.make_addplot(bolling_candles['BoolingDiff'], color='yellow')
    ]

    try:
        save_path = "./pictures/demo/{}.png".format(name)
        mpf.plot(bolling_candles, type='candle', addplot=apds, style='default', volume=True, savefig=save_path)
        # 绘制并推送图片到钉钉
        message = []
        pic_path = picture_upload.upload(save_path)
        message.append("机会出现!")
        message.append("{} {},行业 [{}]".format(name,stock_code, industry))

        call_dingtalk(message, 'news', picurl=pic_path)
    except:
        call_dingtalk("{} 机会出现".format(name),'text')
    send_email("{} {} {}机会出现!".format(name, industry,code), subject="{} 机会出现".format(name))

# SKDJ + MACD
def skdj_macd(stock_frame):

    # 计算macd_skdj
    # 应用MACD和启动点检测
    stock_frame = calculate_macd(stock_frame)
    stock_frame = macd_start_point(stock_frame)
    # 如果当天出现启动点, 返回true
    if stock_frame.at[stock_frame.index[-1], '启动点'] == 1:
        return True
    else:
        return False


def sma(data, time_period):
    return data.rolling(window=time_period).mean()


# 定义 LLV 最低值函数
def llv(data, time_period):
    return data.rolling(window=time_period).min()


# 定义 HHV 最高值函数
def hhv(data, time_period):
    return data.rolling(window=time_period).max()


# 定义计算低价高波信号的函数
def calculate_low_price_high_wave_signal(df, TIME_IV=10):
    # 计算VARC, VARD, VARE, VARF
    df['VARC'] = df['low']
    df['VARD'] = df['low'].shift(1)
    df['VARE'] = (df['VARC'] - df['VARD']).abs().rolling(window=3).mean() / (
                (df['VARC'] - df['VARD']).clip(lower=0).rolling(window=3).mean() + 1e-10) * 100
    df['VARF'] = df['VARE'] / 10.0

    # 计算VAR10和VAR11
    df['VAR10'] = df['low'].rolling(window=TIME_IV).min()
    df['VAR11'] = df['VARF'].rolling(window=TIME_IV).max()

    # 计算低价高波
    df['低价高波'] = df.apply(lambda row: (row['VARF'] + row['VAR11'] * 2) / 2 if row['low'] <= row['VAR10'] else 0,
                              axis=1)
    df['低价高波'] = df['低价高波'].rolling(window=3).mean()

    # 计算低价高波pre和低价高波pre2
    df['低价高波pre'] = df['低价高波'].shift(1)
    df['低价高波pre2'] = df['低价高波'].shift(2)

    # 计算今日、昨日和五周期内信号
    df['今日'] = (df['低价高波pre'] <= 0) & (df['低价高波'] > 0)
    df['昨日'] = (df['低价高波pre2'] <= 0) & (df['低价高波pre'] > 0)
    df['五周期内信号'] = df[['低价高波', '低价高波pre', '低价高波pre2']].apply(lambda x: 1 if any(x > 0) else 0, axis=1)

    # 添加signal列
    df['signal'] = df['今日'] | df['昨日'] | df['五周期内信号']

    # 最后，你可以选择删除中间计算列，如果它们不再需要
    df.drop(['VARC', 'VARD', 'VARE', 'VARF', 'VAR10', 'VAR11', '低价高波pre', '低价高波pre2', '今日', '昨日',
             '五周期内信号'], axis=1, inplace=True)
    # 如果当天出现启动点, 返回true
    return df.at[df.index[-1], 'signal']


# 布林策略 TODO
def boll_select(stock_frame):
    bolling_candles = Boll.Boll(stock_frame, 20, 2)
    # 当天收阳线, 且最近三根日线中存在最低价格下穿布林下轨

    if bolling_candles.at[bolling_candles.index[-1],'Low'] < bolling_candles.at[bolling_candles.index[-1],'LowerBand'] \
        or bolling_candles.at[bolling_candles.index[-2],'Low'] < bolling_candles.at[bolling_candles.index[-2],'LowerBand']\
        or bolling_candles.at[bolling_candles.index[-3],'Low'] < bolling_candles.at[bolling_candles.index[-3],'LowerBand']:
            return True
    return False

# 布林diff为正, 且前三根出现上穿布林中轨, 且当日中轨的上涨百分比高于昨日至少2%, 返回True
def boll_diff(stock_frame):
    bolling_candles = Boll.Boll(stock_frame, 20, 2)
    # if bolling_candles.at[bolling_candles.index[-1],'BoolingDiff'] > 0 :
    #     if (bolling_candles.at[bolling_candles.index[-1],'Close'] > bolling_candles.at[bolling_candles.index[-1],'MiddleBand'])\
    #         and (bolling_candles.at[bolling_candles.index[-1],'Low'] < bolling_candles.at[bolling_candles.index[-1],'MiddleBand']
    #              or bolling_candles.at[bolling_candles.index[-2],'Low'] < bolling_candles.at[bolling_candles.index[-2],'MiddleBand']
    #                 or bolling_candles.at[bolling_candles.index[-3],'Low'] < bolling_candles.at[bolling_candles.index[-3],'MiddleBand'])\
    #         and ((bolling_candles.at[bolling_candles.index[-1],'MiddleBand'] / bolling_candles.at[bolling_candles.index[-2],'MiddleBand']) /
    #               (bolling_candles.at[bolling_candles.index[-2],'MiddleBand'] / bolling_candles.at[bolling_candles.index[-3],'MiddleBand']) > 1.02):
    #             return True

    if bolling_candles.at[bolling_candles.index[-1],'BoolingDiff'] > 0 :
        if (bolling_candles.at[bolling_candles.index[-1],'Close'] > bolling_candles.at[bolling_candles.index[-1],'MiddleBand']):
            if((bolling_candles.at[bolling_candles.index[-1],'MiddleBand'] - bolling_candles.at[bolling_candles.index[-2],'MiddleBand']) /
                  (bolling_candles.at[bolling_candles.index[-2],'MiddleBand'] - bolling_candles.at[bolling_candles.index[-3],'MiddleBand']) > 1.01):
                return True

    return False

def check_volatility_and_middle_band(df, percentile=10):
    """
    判断股票价格的波动率是否处于历史低位，且布林中轨的斜率为正。

    参数：
    df : DataFrame
        包含 'open', 'close', 'high', 'low' 列的股票数据。
    percentile : int, 默认10
        历史低波动率的百分位阈值。

    返回：
    bool
        当且仅当波动率处于历史低位且布林中轨斜率为正时返回True，否则返回False。
    """
    # 检查必要的列是否存在
    required_columns = ['open', 'close', 'high', 'low']
    if not all(col in df.columns for col in required_columns):
        raise ValueError("DataFrame 必须包含 'open', 'close', 'high', 'low' 列。")

    if df.empty:
        return False

    # 计算每日相对波动率：(high - low)/open
    df = df.copy()  # 避免修改原始DataFrame
    df['daily_vol'] = (df['high'] - df['low']) / df['open']

    # 当前波动率（最后一行）
    current_vol = df['daily_vol'].iloc[-1]

    # 历史波动率序列（排除可能的NaN或inf）
    historical_vols = df['daily_vol'].replace([np.inf, -np.inf], np.nan).dropna()

    # 检查是否有足够的历史数据
    if len(historical_vols) == 0:
        condition1 = False
    else:
        # 计算指定百分位数值
        vol_threshold = np.percentile(historical_vols, percentile)
        condition1 = current_vol <= vol_threshold

    # 计算布林中轨（20日SMA）
    window = 20
    df['middle_band'] = df['close'].rolling(window=window).mean()
    middle_band_series = df['middle_band'].dropna()

    # 判断中轨斜率是否为正（最后两个值比较）
    if len(middle_band_series) < 2:
        condition2 = False
    else:
        last_two = middle_band_series.iloc[-2:]
        condition2 = last_two.iloc[-1] > last_two.iloc[-2]

    return condition1 and condition2


import pandas as pd
import numpy as np


def select_high_volatility_stocks(df, percentile=20, n_days = 5):
    """
    筛选当日波动率显著高于历史波动率20百分位，且前一周波动率均低于历史波动率20百分位的股票。

    参数：
    df : DataFrame
        包含 'open', 'high', 'low' 列的股票数据。
    percentile : int, 默认20
        历史波动率的百分位阈值。

    返回：
    DataFrame
        包含符合条件的行的DataFrame。
    """
    required_columns = ['open', 'high', 'low']
    if not all(col in df.columns for col in required_columns):
        raise ValueError("DataFrame必须包含 'open', 'high', 'low' 列")

    if df.empty:
        return pd.DataFrame()

    df = df.copy()

    # 计算每日相对波动率
    df['daily_vol'] = (df['high'] - df['low']) / df['open']

    # 生成历史波动率阈值列（到前一日为止的数据）
    thresholds = []
    for i in range(len(df)):
        if i == 0:
            thresholds.append(np.nan)
            continue
        historical_vols = df['daily_vol'].iloc[:i]
        if len(historical_vols) < 1:
            thresholds.append(np.nan)
        else:
            thresholds.append(np.percentile(historical_vols, percentile))
    df['threshold'] = thresholds

    # 定义条件检查函数
    def check_conditions(row):
        idx = row.name
        if idx < 5 or pd.isna(row['threshold']):
            return False

        # 当日波动率需高于阈值
        cond_today = row['daily_vol'] > row['threshold']

        # 前一周数据索引范围（5个交易日）
        start, end = idx - 5, idx - 1

        # 检查前一周数据是否均低于各自阈值
        prev_vols = df.iloc[start:end + 1]['daily_vol']
        prev_thresholds = df.iloc[start:end + 1]['threshold']

        # 排除含NaN值的行
        valid_check = (
                (prev_vols < prev_thresholds) &
                ~prev_vols.isna() &
                ~prev_thresholds.isna()
        )
        cond_prev_week = valid_check.all()

        return cond_today and cond_prev_week

    # 应用条件检查
    df['signal'] = df.apply(check_conditions, axis=1)

    # 返回信号为True的行
    # 检查最近n天内是否出现信号
    recent_days = df.tail(n_days)
    signal_in_n_days = recent_days.apply(check_conditions, axis=1).any()

    return signal_in_n_days


if __name__ == '__main__':

    # param test's value is False if is True will use simulate trading

    base_dir = r"./"
    n = 3 # 看多少跟K线的最低价低于

    with open(os.path.join(base_dir, "passwords", "passwords.json"), 'r') as f:
        data = json.load(f)

    # 登陆系统
    lg = bs.login()
    # 显示登陆返回信息
    print('login respond error_code:' + lg.error_code)
    print('login respond  error_msg:' + lg.error_msg)

    # 获取A500指数成分股列表
    # 获取中证500成分股
    rs = bs.query_zz500_stocks()
    print('query_zz500 error_code:' + rs.error_code)
    print('query_zz500  error_msg:' + rs.error_msg)

    # 打印结果集
    zz500_stocks = []
    while (rs.error_code == '0') & rs.next():
        # 获取一条记录，将记录合并在一起
        zz500_stocks.append(rs.get_row_data())

    # 获取全部股票代码
    # 设置代理,代理是全局设置,代理失效后可重新设置。参数:ip,proxy_url
    adata.proxy(is_proxy=True, ip='127.0.0.1:10810')
    res_df = adata.stock.info.all_code()
    # 组合 exchange 和 stock_code 成 "SZ.00001" 形式的列表
    # 创建字典，将exchange和stock_code组合成key，short_name作为value
    combined_map = {f"{row['exchange']}.{row['stock_code']}": row['short_name'] for index, row in res_df.iterrows()}



    # 设置日志存储位置
    LOG_DIR = os.path.join(base_dir, "./log")
    if not os.path.exists(LOG_DIR):
        os.makedirs(LOG_DIR)
    LOG_ADDR = os.path.join(LOG_DIR, 'chooseLogA.log')
    logging.basicConfig(filename=LOG_ADDR,format='%(asctime)s - %(message)s', level=logging.INFO, datefmt='%Y-%m-%d %H:%M:%S')
    logger = logging.getLogger(__name__)
    logger.info("start")
    logger.info(combined_map)

    # 图片保存位置
    pic_dir = os.path.join(base_dir, "./pictures")
    if not os.path.exists(pic_dir):
        os.makedirs(pic_dir)
    all_stock_list = []

    # 持续监控
    # while True:
    for c_stock_code, name in combined_map.items():
        # if stock_code.startswith("300") or stock_code.startswith("688"):
        #     continue
        # elif name in all_stock_list:
        #     continue
        stock_code = c_stock_code.split(".")[1]
        if  name in all_stock_list:
            continue
        res_df = adata.stock.market.get_market(stock_code=stock_code, k_type=1, start_date='2024-06-01')
        try:
            info_df = adata.stock.info.get_industry_sw(stock_code=stock_code)
        except:
            pass
        # # 获取单个公司的60天日线数据
        # rs = bs.query_history_k_data_plus(stock_code,
        #                                   "date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,isST",
        #                                   "2024-05-01",
        #                                   "2024-12-03",
        #                                   "d",
        #                                   "3")
        # result = pd.DataFrame(rs.data, columns=rs.fields)
        # 将结果转换为DataFrame
        # data_list = []
        # while (rs.error_code == '0') & rs.next():
        #     # 获取一条记录，将记录合并在一起
        #     data_list.append(rs.get_row_data())
        # df = pd.DataFrame(data_list, columns=rs.fields)
        if res_df.empty:
            continue

        try:
            value = adata.stock.info.get_stock_shares(stock_code=stock_code).loc[0,'total_shares'] * res_df.at[res_df.index[-1],'close']/100000000
            if value > 500:
                continue
        except:
            pass
        # 获取行业分类数据
        industry = ''
        try:
            rs = bs.query_stock_industry(c_stock_code)
            concept = ",".join(adata.stock.info.get_concept_ths(stock_code=stock_code)['name'])

            industry = rs.data[0][3] + concept
        except:
            industry = ""
        # choosed = skdj_macd(res_df)
        # choosed = (calculate_low_price_high_wave_signal(res_df,30) and skdj_macd(res_df))
        # choosed = (boll_diff(res_df) and skdj_macd(res_df))
        choosed = select_high_volatility_stocks(res_df,10,2)
        # select_high_volatility_stocks
        # choosed = boll_select(res_df)
        if choosed:
            all_stock_list.append(name)
            send_mail_and_message(res_df,industry,c_stock_code)
        print(name,choosed,"\n")

    print(all_stock_list)
